跳转至

文章背景与核心概要

在现代深度学习中，如何将高层数学运算高效映射到底层物理硬件，是决定大规模模型训练与推理性能的关键所在。尽管高层张量框架提供了灵活的抽象，但其执行模型往往缺乏最大化硬件利用率所需的全局图可见性，导致复杂的运算不得不严重依赖不透明的手写算子库。

为了弥合这一鸿沟，本文介绍了 Nova 的最新迭代版本——一个自动化的端到端即时（JIT）编译器。它通过直接从计算结构中合成细粒度内核，实现了对硬件映射的绝对控制。该研究将 Nova 的编译流水线扩展为原生支持完整的 Transformer 架构；通过捕获动态执行（eager executions）并将前向与反向传播统一到一个单一的值语义方言（value-semantic dialect）中，Nova 实现了极致的全局图优化。

与依赖僵化的预编译库调用不同，Nova 专注于广泛的跨算子融合，将复杂的因果注意力子图、逐元素操作以及受内存限制的归一化操作直接折叠进单一的融合内核中，从而大幅减少全局内存的往返开销。在 Ada 6000 GPU 上训练完整 GPT-2 架构的评估中，Nova 展现出了卓越的端到端吞吐量，平均达到了每秒 441K 个 token，优于标准动态执行的 406K 以及 torch.compile 的 405K，同时严格保持了数值等价性。

Nova: An End-to-End MLIR Compiler for Deep Learning

Summary

Nova is an automated end-to-end Just-In-Time (JIT) compiler designed to optimize the performance of deep learning models at scale. By synthesizing fine-grained kernels directly from computational structures using MLIR, Nova bridges the gap between high-level tensor abstractions and low-level physical hardware.

Key innovations and results include: * Full Transformer Support: Extends the compilation pipeline to natively capture eager executions and unify forward and backward passes into a single value-semantic dialect. * Compiler-Native Fusion: Eliminates reliance on rigid pre-compiled library calls (such as those for Attention) by aggressively fusing causal attention sub-graphs, element-wise operations, and memory-bound normalizations into single kernels. * Performance: Achieves an average end-to-end throughput of 441K tokens/second when training a full GPT-2 architecture on Ada 6000 GPUs, outperforming standard eager execution (406K tokens/sec) and torch.compile (405K tokens/sec) while strictly preserving numerical parity.

Metadata

  • arXiv ID: arXiv:2608.00029 [cs.AI]
  • Subjects: Artificial Intelligence (cs.AI), Hardware Architecture (cs.AR), Machine Learning (cs.LG), Programming Languages (cs.PL)
  • Authors: Adwaid Suresh, Aparna A, Harshini V M, Jona Delcy C A, Killi Uma Maheswara Rao, Ram Charan Golla, Surendra Vendra
  • Submission Timeline:
  • Submitted: 15 July 2026
  • Last Revised: 2 September 2026 (v3)
  • License: Creative Commons Attribution-NonCommercial-NoDerivatives 4.0 International (View license image: license icon)

Abstract

The performance of deep learning models at scale relies heavily on how effectively high-level mathematical operations are mapped to underlying physical hardware. While high-level tensor frameworks provide flexible abstractions, their execution models inherently lack the whole-graph visibility required to maximize hardware utilization, often forcing a reliance on opaque, hand-written kernel libraries for complex operations like Attention.

To bridge this gap, we present the next iteration of Nova, an automated end-to-end JIT compiler that achieves absolute control over hardware mapping by synthesizing fine-grained kernels directly from the computation's structure. In this work, we extend Nova's compilation pipeline to natively support full Transformer architectures. By capturing eager executions and unifying forward and backward passes into a single value-semantic dialect, Nova unlocks aggressive whole-graph optimizations.

Rather than relying on rigid, pre-compiled library calls, Nova focuses on extensive cross-operator fusions, collapsing complex causal attention sub-graphs, element-wise operations, and memory-bound normalizations directly into single fused kernels to drastically reduce global memory roundtrips. In our evaluations training a full GPT-2 architecture on Ada 6000 GPUs, Nova demonstrates superior end-to-end throughput, averaging 441K tokens/second compared to 406K for our own eager execution and 405K for torch.compile. By drastically reducing memory-bound overheads through compiler-native fusion, Nova enables efficient full LLM compilation on modern hardware while strictly maintaining numerical parity.